"""
Code viewer for interactive exploration.

Provides code viewing capabilities with syntax highlighting,
function extraction, and context display for manual exploration.
"""

import re
from typing import Dict, Any, List, Optional, Tuple
from pathlib import Path

from ...core.base_mcp import CallTreeNode
from ..utils.colors import ColorPrinter, Color


class CodeViewer:
    """Handles code viewing and display for exploration sessions."""
    
    def __init__(self, syntax_highlighting: bool = True):
        self.syntax_highlighting = syntax_highlighting
        self.printer = ColorPrinter()
        self.language_extractors = {
            'python': self._extract_python_function,
            'cpp': self._extract_cpp_function,
            'fortran': self._extract_fortran_subroutine,
            'typescript': self._extract_typescript_function,
            'latex': self._extract_latex_section,
            'markdown': self._extract_markdown_section
        }
    
    async def display_node_code(self, node: CallTreeNode):
        """Display the code for a given node."""
        try:
            self._print_code_header(node)
            
            if node.language in ['latex', 'markdown']:
                # Document content
                content = await self._extract_section_content(node)
                self._display_text_content(content, node.language)
            else:
                # Function code
                code = await self._extract_function_code(node)
                self._display_code_with_highlighting(code, node.language)
            
            self._print_code_footer()
            
        except Exception as e:
            self.printer.print_error(f"Error displaying code: {e}")
    
    def _print_code_header(self, node: CallTreeNode):
        """Print code section header."""
        self.printer.print_separator("─", 60)
        header_text = f"Code for: {node.name}"
        if hasattr(node, 'file_path'):
            header_text += f" ({Path(node.file_path).name}:{node.line_number})"
        
        self.printer.print_colored(header_text, Color.BRIGHT_CYAN)
        self.printer.print_separator("─", 60)
    
    def _print_code_footer(self):
        """Print code section footer."""
        self.printer.print_separator("─", 60)
    
    async def _extract_function_code(self, node: CallTreeNode) -> str:
        """Extract function code from source file."""
        try:
            file_path = Path(node.file_path)
            if not file_path.exists():
                return f"File not found: {node.file_path}"
            
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # Use language-specific extractor
            extractor = self.language_extractors.get(node.language)
            if extractor:
                return extractor(lines, node)
            else:
                return self._extract_generic_function(lines, node)
                
        except Exception as e:
            return f"Error reading source: {e}"
    
    async def _extract_section_content(self, node: CallTreeNode) -> str:
        """Extract document section content."""
        try:
            file_path = Path(node.file_path)
            if not file_path.exists():
                return f"File not found: {node.file_path}"
            
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            if node.language == 'latex':
                return self._extract_latex_section_content(content, node)
            elif node.language == 'markdown':
                return self._extract_markdown_section_content(content, node)
            else:
                return content[:1000] + "..." if len(content) > 1000 else content
                
        except Exception as e:
            return f"Error reading document: {e}"
    
    def _extract_python_function(self, lines: List[str], node: CallTreeNode) -> str:
        """Extract Python function with proper indentation."""
        start_line = max(0, node.line_number - 1)  # Convert to 0-based index
        
        # Find function definition
        function_pattern = re.compile(r'^\s*(async\s+)?def\s+' + re.escape(node.name) + r'\s*\(')
        
        for i in range(start_line, len(lines)):
            if function_pattern.search(lines[i]):
                start_line = i
                break
        else:
            # Function not found at expected line, search nearby
            for i in range(max(0, start_line - 5), min(len(lines), start_line + 10)):
                if function_pattern.search(lines[i]):
                    start_line = i
                    break
            else:
                return f"Function '{node.name}' not found around line {node.line_number}"
        
        # Extract function with proper indentation
        function_lines = [lines[start_line]]
        base_indent = len(lines[start_line]) - len(lines[start_line].lstrip())
        
        for i in range(start_line + 1, len(lines)):
            line = lines[i]
            
            # Empty line
            if line.strip() == "":
                function_lines.append(line)
                continue
            
            # Check indentation to determine end of function
            line_indent = len(line) - len(line.lstrip())
            if line_indent <= base_indent and line.strip():
                # Check if this is a decorator for next function
                if line.strip().startswith('@'):
                    break
                # Check if this is another function/class definition
                if re.match(r'^\s*(def|class|async\s+def)\s+', line):
                    break
                # If it's a statement at module level, function ends
                if line_indent == 0:
                    break
            
            function_lines.append(line)
            
            # Limit function length for display
            if len(function_lines) > 100:
                function_lines.append("    # ... (function continues) ...\n")
                break
        
        return ''.join(function_lines)
    
    def _extract_cpp_function(self, lines: List[str], node: CallTreeNode) -> str:
        """Extract C++ function code."""
        start_line = max(0, node.line_number - 1)
        function_name = node.name
        
        # Look for function definition
        patterns = [
            rf'\b{re.escape(function_name)}\s*\(',  # Function call/definition
            rf'\b\w+\s+{re.escape(function_name)}\s*\(',  # Return type + function
            rf'\b\w+::{re.escape(function_name)}\s*\('  # Class member function
        ]
        
        found_start = None
        for i in range(max(0, start_line - 5), min(len(lines), start_line + 10)):
            line = lines[i].strip()
            for pattern in patterns:
                if re.search(pattern, line):
                    found_start = i
                    break
            if found_start is not None:
                break
        
        if found_start is None:
            return f"Function '{function_name}' not found around line {node.line_number}"
        
        # Extract function body
        function_lines = []
        brace_count = 0
        in_function = False
        
        for i in range(found_start, len(lines)):
            line = lines[i]
            function_lines.append(line)
            
            # Count braces to find function end
            for char in line:
                if char == '{':
                    brace_count += 1
                    in_function = True
                elif char == '}':
                    brace_count -= 1
            
            # Function complete when braces match
            if in_function and brace_count == 0:
                break
            
            # Limit length
            if len(function_lines) > 150:
                function_lines.append("    // ... (function continues) ...\n")
                break
        
        return ''.join(function_lines)
    
    def _extract_fortran_subroutine(self, lines: List[str], node: CallTreeNode) -> str:
        """Extract Fortran subroutine/function code."""
        start_line = max(0, node.line_number - 1)
        name = node.name.lower()  # Fortran is case-insensitive
        
        # Look for subroutine/function definition
        patterns = [
            rf'^\s*subroutine\s+{re.escape(name)}\s*\(',
            rf'^\s*function\s+{re.escape(name)}\s*\(',
            rf'^\s*.*function\s+{re.escape(name)}\s*\('  # typed function
        ]
        
        found_start = None
        for i in range(max(0, start_line - 5), min(len(lines), start_line + 10)):
            line = lines[i].lower()
            for pattern in patterns:
                if re.search(pattern, line):
                    found_start = i
                    break
            if found_start is not None:
                break
        
        if found_start is None:
            return f"Subroutine/Function '{node.name}' not found around line {node.line_number}"
        
        # Extract until end statement
        function_lines = []
        for i in range(found_start, len(lines)):
            line = lines[i]
            function_lines.append(line)
            
            # Check for end statement
            line_lower = line.lower().strip()
            if (line_lower.startswith('end subroutine') or 
                line_lower.startswith('end function') or
                line_lower == 'end'):
                break
            
            # Limit length
            if len(function_lines) > 200:
                function_lines.append("    ! ... (subroutine continues) ...\n")
                break
        
        return ''.join(function_lines)
    
    def _extract_typescript_function(self, lines: List[str], node: CallTreeNode) -> str:
        """Extract TypeScript/JavaScript function code."""
        start_line = max(0, node.line_number - 1)
        function_name = node.name
        
        # Look for function definition
        patterns = [
            rf'function\s+{re.escape(function_name)}\s*\(',
            rf'const\s+{re.escape(function_name)}\s*=',
            rf'let\s+{re.escape(function_name)}\s*=',
            rf'var\s+{re.escape(function_name)}\s*=',
            rf'{re.escape(function_name)}\s*:\s*function',
            rf'{re.escape(function_name)}\s*\('  # Method definition
        ]
        
        found_start = None
        for i in range(max(0, start_line - 5), min(len(lines), start_line + 10)):
            line = lines[i]
            for pattern in patterns:
                if re.search(pattern, line):
                    found_start = i
                    break
            if found_start is not None:
                break
        
        if found_start is None:
            return f"Function '{function_name}' not found around line {node.line_number}"
        
        # Extract function body (similar to C++ but with different syntax)
        function_lines = []
        brace_count = 0
        in_function = False
        
        for i in range(found_start, len(lines)):
            line = lines[i]
            function_lines.append(line)
            
            # Count braces
            for char in line:
                if char == '{':
                    brace_count += 1
                    in_function = True
                elif char == '}':
                    brace_count -= 1
            
            # Function complete
            if in_function and brace_count == 0:
                break
            
            # Limit length
            if len(function_lines) > 120:
                function_lines.append("    // ... (function continues) ...\n")
                break
        
        return ''.join(function_lines)
    
    def _extract_latex_section(self, lines: List[str], node: CallTreeNode) -> str:
        """Extract LaTeX section content."""
        # This would extract the section starting from the node's line
        start_line = max(0, node.line_number - 1)
        
        # Find section start
        section_patterns = [
            r'\\chapter\{',
            r'\\section\{',
            r'\\subsection\{',
            r'\\subsubsection\{'
        ]
        
        found_start = start_line
        for i in range(max(0, start_line - 2), min(len(lines), start_line + 5)):
            line = lines[i]
            for pattern in section_patterns:
                if re.search(pattern, line):
                    found_start = i
                    break
        
        # Extract section content (simplified)
        section_lines = []
        for i in range(found_start, min(len(lines), found_start + 50)):
            line = lines[i]
            section_lines.append(line)
            
            # Stop at next section of same or higher level
            if i > found_start and any(re.search(pattern, line) for pattern in section_patterns):
                break
        
        return ''.join(section_lines)
    
    def _extract_markdown_section(self, lines: List[str], node: CallTreeNode) -> str:
        """Extract Markdown section content."""
        start_line = max(0, node.line_number - 1)
        
        # Find header
        header_line = None
        for i in range(max(0, start_line - 2), min(len(lines), start_line + 5)):
            if lines[i].strip().startswith('#'):
                header_line = i
                break
        
        if header_line is None:
            header_line = start_line
        
        # Extract content until next header of same or higher level
        section_lines = []
        header_level = len(lines[header_line].split()[0])  # Count # characters
        
        for i in range(header_line, min(len(lines), header_line + 100)):
            line = lines[i]
            section_lines.append(line)
            
            # Stop at next header of same or higher level
            if (i > header_line and 
                line.strip().startswith('#') and 
                len(line.split()[0]) <= header_level):
                break
        
        return ''.join(section_lines)
    
    def _extract_generic_function(self, lines: List[str], node: CallTreeNode) -> str:
        """Generic function extraction for unknown languages."""
        start_line = max(0, node.line_number - 1)
        end_line = min(len(lines), start_line + 50)  # Show 50 lines max
        
        extracted_lines = lines[start_line:end_line]
        
        if len(extracted_lines) == 50:
            extracted_lines.append("... (truncated) ...\n")
        
        return ''.join(extracted_lines)
    
    def _display_code_with_highlighting(self, code: str, language: str):
        """Display code with syntax highlighting."""
        if not code.strip():
            self.printer.print_warning("No code content available")
            return
        
        if self.syntax_highlighting:
            try:
                # Try to use pygments for syntax highlighting
                from pygments import highlight
                from pygments.lexers import get_lexer_by_name
                from pygments.formatters import TerminalFormatter
                
                # Map our language names to pygments lexer names
                lexer_map = {
                    'python': 'python',
                    'cpp': 'cpp',
                    'fortran': 'fortran',
                    'typescript': 'typescript',
                    'javascript': 'javascript'
                }
                
                lexer_name = lexer_map.get(language, language)
                lexer = get_lexer_by_name(lexer_name)
                formatter = TerminalFormatter()
                highlighted = highlight(code, lexer, formatter)
                print(highlighted.rstrip())
                return
                
            except ImportError:
                # Pygments not available, fall back to basic highlighting
                pass
            except Exception:
                # Any other error, fall back to basic highlighting
                pass
        
        # Fallback to basic syntax highlighting
        self._display_code_with_basic_highlighting(code, language)
    
    def _display_code_with_basic_highlighting(self, code: str, language: str):
        """Display code with basic color highlighting."""
        lines = code.split('\n')
        
        for line in lines:
            stripped = line.strip()
            
            # Comments
            if self._is_comment_line(stripped, language):
                self.printer.print_colored(line, Color.BRIGHT_BLACK)
            # Keywords/definitions
            elif self._contains_keywords(stripped, language):
                self.printer.print_colored(line, Color.BRIGHT_BLUE)
            # String literals
            elif self._contains_strings(stripped):
                self.printer.print_colored(line, Color.BRIGHT_GREEN)
            # Preprocessor/import statements
            elif self._is_preprocessor_line(stripped, language):
                self.printer.print_colored(line, Color.BRIGHT_MAGENTA)
            # Regular code
            else:
                self.printer.print_colored(line, Color.WHITE)
    
    def _is_comment_line(self, line: str, language: str) -> bool:
        """Check if line is a comment."""
        comment_chars = {
            'python': '#',
            'cpp': '//',
            'fortran': '!',
            'typescript': '//',
            'latex': '%',
            'markdown': None
        }
        
        char = comment_chars.get(language)
        return char and line.startswith(char)
    
    def _contains_keywords(self, line: str, language: str) -> bool:
        """Check if line contains language keywords."""
        keywords = {
            'python': ['def ', 'class ', 'if ', 'for ', 'while ', 'import ', 'from '],
            'cpp': ['int ', 'void ', 'class ', 'struct ', 'if ', 'for ', 'while '],
            'fortran': ['subroutine ', 'function ', 'if ', 'do ', 'program '],
            'typescript': ['function ', 'class ', 'interface ', 'if ', 'for ', 'while ']
        }
        
        lang_keywords = keywords.get(language, [])
        return any(keyword in line.lower() for keyword in lang_keywords)
    
    def _contains_strings(self, line: str) -> bool:
        """Check if line contains string literals."""
        return '"' in line or "'" in line
    
    def _is_preprocessor_line(self, line: str, language: str) -> bool:
        """Check if line is a preprocessor directive."""
        if language == 'cpp':
            return line.startswith('#include') or line.startswith('#define')
        elif language == 'python':
            return line.startswith('import ') or line.startswith('from ')
        elif language == 'fortran':
            return line.lower().startswith('use ')
        elif language == 'typescript':
            return line.startswith('import ') or line.startswith('export ')
        
        return False
    
    def _display_text_content(self, content: str, doc_type: str):
        """Display document content with basic formatting."""
        if not content.strip():
            self.printer.print_warning("No content available")
            return
        
        lines = content.split('\n')
        
        for line in lines:
            stripped = line.strip()
            
            if doc_type == 'latex':
                # LaTeX-specific highlighting
                if stripped.startswith('\\chapter') or stripped.startswith('\\section'):
                    self.printer.print_colored(line, Color.BRIGHT_CYAN)
                elif stripped.startswith('\\'):
                    self.printer.print_colored(line, Color.BRIGHT_YELLOW)
                elif stripped.startswith('%'):
                    self.printer.print_colored(line, Color.BRIGHT_BLACK)
                else:
                    self.printer.print_colored(line, Color.WHITE)
                    
            elif doc_type == 'markdown':
                # Markdown-specific highlighting
                if stripped.startswith('#'):
                    self.printer.print_colored(line, Color.BRIGHT_CYAN)
                elif stripped.startswith('```'):
                    self.printer.print_colored(line, Color.BRIGHT_YELLOW)
                elif stripped.startswith('-') or stripped.startswith('*'):
                    self.printer.print_colored(line, Color.BRIGHT_GREEN)
                else:
                    self.printer.print_colored(line, Color.WHITE)
            else:
                self.printer.print_colored(line, Color.WHITE)
    
    def _extract_latex_section_content(self, content: str, node: CallTreeNode) -> str:
        """Extract specific LaTeX section content."""
        # This would implement more sophisticated LaTeX parsing
        # For now, return a portion around the node's line
        lines = content.split('\n')
        start_idx = max(0, node.line_number - 1)
        end_idx = min(len(lines), start_idx + 50)
        
        return '\n'.join(lines[start_idx:end_idx])
    
    def _extract_markdown_section_content(self, content: str, node: CallTreeNode) -> str:
        """Extract specific Markdown section content."""
        # Similar to LaTeX but for Markdown
        lines = content.split('\n')
        start_idx = max(0, node.line_number - 1)
        end_idx = min(len(lines), start_idx + 50)
        
        return '\n'.join(lines[start_idx:end_idx])